import json
import os

def labelme2yolo(json_dir, labels_dir, class_names):
    if not os.path.exists(labels_dir):
        os.makedirs(labels_dir)

    for json_file in os.listdir(json_dir):
        if json_file.endswith('.json'):
            json_path = os.path.join(json_dir, json_file)
            with open(json_path, 'r') as f:
                data = json.load(f)

            image_width = data['imageWidth']
            image_height = data['imageHeight']
            shapes = data['shapes']

            txt_filename = os.path.basename(json_file).replace('.json', '.txt')
            txt_path = os.path.join(labels_dir, txt_filename)

            with open(txt_path, 'w') as file_handle:
                for shape in shapes:
                    label = shape['label']
                    if label not in class_names:
                        continue
                    class_id = class_names.index(label)
                    points = shape['points']

                    if shape['shape_type'] == 'rectangle':
                        x1, y1 = points[0]
                        x2, y2 = points[1]
                    elif shape['shape_type'] == 'polygon':
                        x1 = min(point[0] for point in points)
                        y1 = min(point[1] for point in points)
                        x2 = max(point[0] for point in points)
                        y2 = max(point[1] for point in points)
                    else:
                        continue

                    x_center = (x1 + x2) / 2.0 / image_width
                    y_center = (y1 + y2) / 2.0 / image_height
                    width = (x2 - x1) / image_width
                    height = (y2 - y1) / image_height

                    line_content = [str(class_id), str(x_center), str(y_center), str(width), str(height)]
                    file_handle.write(" ".join(line_content) + "\n")

            print(f"转换完成：{txt_filename}")

# 示例用法
json_dir = 'F:/data/pdf-image-15pc'
labels_dir = 'F:/data/pdf-image-15pc'
class_names = ['txtblock','table','titlebar']  # 你的类别名称列表

labelme2yolo(json_dir, labels_dir, class_names)